from typing import Dict, List

import numpy as np
from pydantic import BaseModel
from volcenginesdkarkruntime import Ark

REGION_CN_NORTH1 = "cn-north-1"
HOST = "api-vikingdb.volces.com"
EMBEDDING_V2_PATH = "/api/data/embedding/version/2"
SERVICE = "air"


class EmbModel(object):
    def __init__(self, model_name, params=None):
        self._model_name = model_name
        self._params = params

    @property
    def model_name(self):
        return self._model_name

    @property
    def params(self):
        return self._params


class RawData(object):
    def __init__(self, data_type, text="", image=""):
        self._data_type = data_type
        self._text = text
        self._image = image

    @property
    def data_type(self):
        return self._data_type

    @property
    def text(self):
        return self._text

    @property
    def image(self):
        return self._image


class Credentials(object):
    def __init__(self, ak, sk, service, region, session_token=''):
        self.ak = ak
        self.sk = sk
        self.service = service
        self.region = region
        self.session_token = session_token

    def set_ak(self, ak):
        self.ak = ak

    def set_sk(self, sk):
        self.sk = sk

    def set_session_token(self, session_token):
        self.session_token = session_token


class ArkException(Exception):
    def __init__(self, code, request_id, message=None):
        self.code = code
        self.request_id = request_id
        if message is not None:
            self.message = message
        else:
            self.message = "unknown error, please contact customer service, request_id:{}".format(self.request_id)

    def __str__(self):
        return self.message


class EmbeddingV2Response(BaseModel):
    code: str
    message: str
    request_id: str
    data: dict


def sliced_norm_l2(vec: List[float], dim=2048) -> List[float]:
    # dim 取值 512,1024,2048
    norm = float(np.linalg.norm(vec[:dim]))
    return [v / norm for v in vec[:dim]]


class ArkFactory:
    _client_dict: Dict[str, Ark] = {}

    @staticmethod
    def get_instance(ak, sk):
        if ak not in ArkFactory._client_dict:
            ArkFactory._client_dict[ak] = Ark(ak=ak, sk=sk)
        return ArkFactory._client_dict[ak]
